# -*- coding: utf-8 -*-
"""
基于 Flask + YOLOv8 的目标检测服务
功能：
1. 自动下载模型
2. 提供 /detect 接口：上传图片 -> 返回检测结果 + 画框图
3. 提供 /history 接口：查询历史预测记录
4. 提供 /clear_history 接口：清空历史记录
"""

import io
import sqlite3
from datetime import datetime
import os
import requests
from flask import Flask, request, jsonify
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO
from filelock import FileLock
from tqdm import tqdm

# -------------------- 初始化 --------------------
app = Flask(__name__)

# -------------------- 模型下载 --------------------
MODEL_PATH = "best.pt"
MODEL_URL = "https://yesongyun-ld-1311137168.cos.ap-shanghai.myqcloud.com/best.pt"

# 如果本地不存在模型，则从云端下载
if not os.path.isfile(MODEL_PATH):
    print(f"[INFO] {MODEL_PATH} 不存在，开始下载……")
    try:
        response = requests.get(MODEL_URL, stream=True, timeout=30)
        response.raise_for_status()
        total_size = int(response.headers.get('content-length', 0))
        # 使用 tqdm 显示下载进度条
        with open(MODEL_PATH, "wb") as f, tqdm(
                desc=MODEL_PATH,
                total=total_size,
                unit='B',
                unit_scale=True,
                unit_divisor=1024
                ) as bar:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
                    bar.update(len(chunk))
        print(f"[INFO] 下载完成：{MODEL_PATH}")
    except Exception as e:
        print(f"[ERROR] 下载失败：{e}")
        exit(1)

# -------------------- 加载模型 --------------------
model = YOLO(MODEL_PATH)

# -------------------- 预热模型 --------------------
# 避免第一次推理延迟过高，先跑一张空图
print("[INFO] Warming up model ...")
warmup_img = Image.new("RGB", (640, 640), color=(114, 114, 114))
_ = model.predict(warmup_img, verbose=False)
print("[INFO] Warmup done.")

# -------------------- 数据库初始化 --------------------
DB_NAME = "predictions.db"
# check_same_thread=False 允许在多线程环境中共用连接
conn = sqlite3.connect(DB_NAME, check_same_thread=False)
cursor = conn.cursor()
# 创建预测结果表（如果不存在）
cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS predictions
        (
            id        INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp TEXT,
            img_name  TEXT,
            cls       INTEGER,
            label     TEXT,
            conf      REAL,
            x1        REAL,
            y1        REAL,
            x2        REAL,
            y2        REAL
        );
        """
        )
conn.commit()


# -------------------- 工具函数 --------------------
def predict_and_annotate(file_bytes: bytes):
    """
    对上传的图片进行推理并画框
    :param file_bytes: 原始图片字节
    :return:
        annotated_bytes: 画框后的图片字节
        cls_list:        检测到的类别名称列表
        records:         结构化记录，用于写入数据库
    """
    # 将字节流转为 PIL 图像
    global label
    img = Image.open(io.BytesIO(file_bytes)).convert("RGB")

    # 使用 YOLO 模型推理
    results = model(img)

    all_boxes = []
    if results[0].boxes is not None:
        for b in results[0].boxes:
            cls = int(b.cls)  # 类别索引
            label = model.names[cls]  # 类别名称
            all_boxes.append(
                    {
                            "label": label,
                            "conf" : float(b.conf),
                            "xyxy" : b.xyxy[0].tolist()  # [x1, y1, x2, y2]
                            }
                    )

    # 在图上画框和文字
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.truetype("arial.ttf", 16)
    except:
        font = ImageFont.load_default()

    for box in all_boxes:
        x1, y1, x2, y2 = box["xyxy"]
        label = box["label"]
        conf = box["conf"]
        text = f"{label} {conf:.2f}"
        draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
        draw.text((x1, y1 - 15), text, fill="red", font=font)

    # 将 PIL 图像转回字节流
    buf = io.BytesIO()
    img.save(buf, format="JPEG")
    annotated_bytes = buf.getvalue()

    # 类别名称列表
    cls_list = [b["label"] for b in all_boxes]

    # 构造数据库记录
    records = [
            {
                    "cls"  : int(b["label"] != label),  # 与原逻辑保持一致，可改为 cls 索引
                    "label": b["label"],
                    "conf" : b["conf"],
                    "x1"   : b["xyxy"][0],
                    "y1"   : b["xyxy"][1],
                    "x2"   : b["xyxy"][2],
                    "y2"   : b["xyxy"][3],
                    }
            for b in all_boxes
            ]

    return annotated_bytes, cls_list, records


# -------------------- 路由 --------------------
@app.route("/detect", methods=["POST"])
def detect():
    """
    检测接口
    表单字段：
        file: 图片文件
    返回：
        cls_list: 检测到的类别名称列表
        img:      画框后的图片字节（十六进制字符串）
    """
    if "file" not in request.files:
        return jsonify({"error": "Missing file"}), 400

    file = request.files["file"]
    img_name = file.filename
    file_bytes = file.read()

    # 推理并画框
    annotated_bytes, cls_list, records = predict_and_annotate(file_bytes)

    # 将结果写入数据库
    timestamp = datetime.now().isoformat(sep=" ", timespec="seconds")
    DB_LOCK = FileLock("predictions.db.lock")
    with DB_LOCK:
        for r in records:
            cursor.execute(
                    "INSERT INTO predictions (timestamp, img_name, cls, label, conf, x1, y1, x2, y2) "
                    "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                    (
                            timestamp, img_name, r["cls"], r["label"], r["conf"],
                            r["x1"], r["y1"], r["x2"], r["y2"]
                            )
                    )
        conn.commit()

    # 将字节流转为十六进制字符串返回，前端可用 base64 解码
    return jsonify(
            {
                    "cls_list": cls_list,
                    "img"     : annotated_bytes.hex()
                    }
            )


@app.route("/history", methods=["GET"])
def history():
    """
    查询历史记录
    可选参数：
        limit: 返回条数，默认 50
    """
    limit = request.args.get("limit", 50, type=int)
    cursor.execute(
            """
            SELECT id,
                   timestamp,
                   img_name,
                   cls,
                   label,
                   conf,
                   x1,
                   y1,
                   x2,
                   y2
            FROM predictions
            ORDER BY id DESC
            LIMIT ?
            """, (limit,)
            )
    rows = cursor.fetchall()
    keys = [
            "id", "timestamp", "img_name", "cls", "label",
            "conf", "x1", "y1", "x2", "y2"
            ]
    return jsonify([dict(zip(keys, row)) for row in rows])


@app.route("/clear_history", methods=["POST"])
def clear_history():
    """
    清空历史记录
    """
    try:
        DB_LOCK = FileLock("predictions.db.lock")
        with DB_LOCK:
            cursor.execute("DELETE FROM predictions")  # 清空表
            cursor.execute("DELETE FROM sqlite_sequence WHERE name='predictions'")  # 重置自增 id
            conn.commit()
        return jsonify({"message": "历史记录已清空"}), 200
    except Exception as e:
        return jsonify({"error": f"清空失败: {str(e)}"}), 500


# -------------------- 启动 --------------------
if __name__ == "__main__":
    # 0.0.0.0 允许外部访问，debug=True 热重载
    app.run(host="0.0.0.0", port=5000, debug=True)
